# -*- coding: utf-8 -*-
"""
Created on Sun Jan  3 17:31:44 2016

@author: liuyi05
"""

import logging
import os
import numpy as np
import pandas as pd
import datetime as dt
from abc import abstractmethod
from core.calc import Calc

class Strategy(object):
    def __init__(self):
        pass

    @abstractmethod
    def gen_op_series(self, df):
        pass
    
    @abstractmethod
    def gen_op(self, df):
        pass


class Turtle(Strategy):
    def __init__(self, min_period, max_period):
        self.min_period = min_period
        self.max_period = max_period
        self.label_min = "min" + str(min_period)
        self.label_max = "max" + str(max_period)
        
    def prepare_cols(self, df):
        if self.label_min not in df:
            Calc.calc_min(df, self.min_period, self.label_min)
        if self.label_max not in df:
            Calc.calc_max(df, self.max_period, self.label_max)

    def gen_op_series(self, df):
        self.prepare_cols(df)
        df["op"] = df.apply(lambda r: 1 if r["high"] > r[self.label_max]
                            else -1 if r["low"] < r[self.label_min]
                            else 0, axis=1)
        
    def gen_op(self, df):
        self.prepare_cols(df)
        row = df.iloc[-1]
        if row["high"] > row[self.label_max]:
            return 1
        elif row["low"] < row[self.label_min]:
            return -1
        else:
            return 0
            
class MovingAverage(Strategy):
    def __init__(self, short_period, long_period):
        self.short_period = short_period
        self.long_period = long_period
        self.label_short = "ma" + str(short_period)
        self.label_long = "ma" + str(long_period)
        
    def prepare_cols(self, df):
        if self.label_short not in df:
            Calc.calc_ma(df, self.short_period, self.label_short)
        if self.label_long not in df:
            Calc.calc_ma(df, self.long_period, self.label_long)

    def gen_op_series(self, df):
        self.prepare_cols(df)
        pos = df[self.label_short] > df[self.label_long]
        df["op"] = 0
        df.ix[(pos == True) & (pos.shift() == False), "op"] = 1
        df.ix[(pos == False) & (pos.shift() == True), "op"] = -1
        
    def gen_op(self, df):
        assert(df.shape[0] > 1)
        self.prepare_cols(df)
        yesterday = df.iloc[-2]
        today = df.iloc[-1]
        if today[self.label_short] > today[self.label_long] \
                and yesterday[self.label_short] < yesterday[self.label_long]:
            return 1
        elif today[self.label_short] < today[self.label_long] \
                and yesterday[self.label_short] > yesterday[self.label_long]:
            return -1
        else:
            return 0


#class KDJ(Strategy):
#    def __init__(self, period):
#        self.period = period
#        self.label_rsv = "kdj_rsv" + str(period)
#        self.label_k = "kdj_k" + str(period)
#        self.label_d = "kdj_d" + str(period)
#        self.label_j = "kdj_j" + str(period)
#        
#    def calc_rsv(self, df):
#        if self.label_rsv in df:
#            return False
#        mi = pd.rolling_min(df["low"], self.period)
#        ma = pd.rolling_max(df["high"], self.period)
#        for i in range(min(self.period - 1, df.shape[0])):
#            mi[i] = min(df["low"][:i+1])
#            ma[i] = max(df["high"][:i+1])
#        eps = 1e-7
#        mi_eps = mi - eps
#        df[self.label_rsv] = (df["close"] - mi_eps) / (ma - mi_eps) * 100
#        df.ix[ma == mi, self.label_rsv] = 50
#        return True
#        
#    def calc_kdj(self, df):
#        size = df.shape[0]
#        assert(size > 0)
#        if self.label_k in df and self.label_d in df and self.label_j in df:
#            return False
#        if self.label_rsv not in df:
#            self.calc_rsv(df)
#        rsv = df[self.label_rsv]
#        kdj_k = pd.Series(np.zeros(size))
#        kdj_d = pd.Series(np.zeros(size))
#        kdj_j = pd.Series(np.zeros(size))
#        kdj_k[0] = (100 + rsv[0]) / 3.0
#        kdj_d[0] = (100 + kdj_k[0]) / 3.0
#        kdj_j[0] = 3.0 * kdj_k[0] - 2.0 * kdj_d[0]
#        for i in range(1, size):
#            kdj_k[i] = (2 * kdj_k[i-1] + rsv[i]) / 3.0
#            kdj_d[i] = (2 * kdj_d[i-1] + kdj_k[i]) / 3.0
#            kdj_j[i] = 3.0 * kdj_k[i] - 2.0 * kdj_d[i]
#        df[self.label_k] = kdj_k
#        df[self.label_d] = kdj_d
#        df[self.label_j] = kdj_j
#        return True
#        
#    def prepare_cols(self, df):
#        if self.label_k not in df or self.label_d not in df \
#                or self.label_j not in df:
#            self.calc_kdj(df)
#
#    def gen_op_series(self, df):
#        self.prepare_cols(df)
#        pass
#        
#    def gen_op(self, df):
#        self.prepare_cols(df)
#        row = df.iloc[-1]
#        return 0

def latest_workday():
    oneday = dt.timedelta(days=1)
    date = dt.datetime.today()
    hour = dt.datetime.now().hour
    if hour < 16:
        date -= oneday
    while date.weekday() in (5, 6):
        date -= oneday
    return date
    
class Record:
    def __init__(self, code, name, industry, turnover, op):
        self.code = code
        self.name = name
        self.industry = industry
        self.turnover = turnover
        self.op = op
    
    def __str__(self):
        return "%s %s [%s] %.2f%% %s" % (self.code, self.name, self.industry, 
                                         self.turnover, self.op)
        
def main():
    code_list_path = "data/code.list"
    basic_info_path = "data/stock_basic_list.csv"
    #info = pd.read_csv(basic_info_path, encoding="gbk",
    #                   dtype=str)[["code", "name", "industry"]]
    info = pd.read_csv(basic_info_path, encoding="utf8",
                       dtype=str)[["code", "name", "industry"]]
    info.set_index("code", inplace=True)
    dir_path = "data/latest/"
    
    strategy = Turtle(10, 20)
    #strategy = MovingAverage(20, 60)
    
    date = latest_workday().strftime("%Y-%m-%d")
    #date = "2016-07-21"
    print("check date:", date)
    with open(code_list_path) as file:
        code_list = [line.rstrip() for line in file.readlines()]
        buy_list = []
        sell_list = []
        print("code name [industry] turnover operation")
        for code in code_list:
            path = dir_path + code + ".asc"
            if not os.path.exists(path):
                continue
            df = pd.read_csv(path, index_col="date")
            if df.index[-1] != date:
                continue
            op = strategy.gen_op(df)
            if op == 1:
                buy_list.append(Record(code, info.loc[code, "name"], 
                                 info.loc[code, "industry"], 
                                 df.iloc[-1]["turnover"], "buy"))
            if op == -1:
                sell_list.append(Record(code, info.loc[code, "name"], 
                                  info.loc[code, "industry"], 
                                  df.iloc[-1]["turnover"], "sell"))
        buy_list.sort(key=lambda x: (x.industry, x.turnover), reverse=True)
        for record in buy_list:
            print(record)
#        print()
#        sell_list.sort(key=lambda x: (x[2], x[3]))
#        for record in sell_list:
#            print(record)
        
        attention = set(["600109", "600276", "002230", "002310", "002614",
                         "300024", "300251", "002340", "000413"
                         "000157","600519","000568","601166","600547","300332",
                         "000538","601699","600970","300027","000338","300146",
                         "300002","300104","300315","300359","300364",
                         "600030","603019","601186","601766","002050","300294"])
                         
        print("\nWatch List:")
        for record in buy_list:
            if record.code in attention:
                print(record)
                
        for record in sell_list:
            if record.code in attention:
                print(record)

if __name__ == "__main__":
    logging.basicConfig(filename="stock.log")
    logging.getLogger().addHandler(logging.StreamHandler())
    main()